
#include "gh.h"
#include "gh.ptx.h"
#include "gh.cu.h"

static const char *nVertexProgramSrc =
"!!NVvp5.0"
NL
NL "main:"
NL "    MOV.F32 result.position, vertex.attrib[0];"
NL "    RET;"
NL "END";

static const char *nFragmentProgramSrc =
"!!NVfp5.0"
NL
NL "main:"
NL "    TEX result.color.rgb, fragment.position, texture[0], RECT;"
NL "    MOV.F32 result.color.a, 1.0;"
NL "    RET;"
NL "END";

static const char *nTriModuleSrc =
NL ".version        2.3"
NL ".target         sm_20"
NL ".address_size   32"

NL ".global .surfref    nFb;"
NL ".const .v2 .f32     nPoint[3] = { { -0.7, -0.7 }, { 0.7, -0.7 }, { 0.0, 0.7 } };"
NL ".const .v2 .f32     nRotZ[2] = { { 1.0, 0.0 }, { 0.0, 1.0 } };"

NL ".func (.reg .f32 sign) ComputeSign("
NL "    .reg .v2 .f32   p0,"
NL "    .reg .v2 .f32   p1,"
NL "    .reg .v2 .f32   p"
NL ") {"
NL "    .reg .v2 .f32   AB;"
NL "    .reg .f32       C;"
NL "    sub.f32         AB.x, p1.y, p0.y;"
NL "    sub.f32         AB.y, p0.x, p1.x;"
        /* C = p1.x * p0.y - p0.x * p1.y; */
NL "    {"
NL "    .reg .f32   t;"
NL "    neg.f32     t, p0.x;"
NL "    mul.f32     t, t, p1.y;"
NL "    fma.rn.f32  C, p1.x, p0.y, t;"
NL "    }"
        /*  sign = dot(AB, p) + C; */
        dot_v2_f32  ("sign", "AB", "p")
NL "    add.f32     sign, sign, C;"
NL "    ret;"
NL "}"

NL ".entry Draw() {"
        /* compute x coordinate: x = %ctaid.x * %ntid.x + %tid.x; */
NL "    .reg .u32   x, y, blockId, blockSize, threadId;"
NL "    mov.u32     blockId, %ctaid.x;"
NL "    mov.u32     blockSize, %ntid.x;"
NL "    mov.u32     threadId, %tid.x;"
NL "    mad.lo.u32  x, blockId, blockSize, threadId;"
        /* compute y coordinate: y = %ctaid.y * %ntid.y + %tid.y; */
NL "    mov.u32     blockId, %ctaid.y;"
NL "    mov.u32     blockSize, %ntid.y;"
NL "    mov.u32     threadId, %tid.y;"
NL "    mad.lo.u32  y, blockId, blockSize, threadId;"
        /* convert x, y to the [-1.0, 1.0] range */
NL "    .reg .v2 .f32   p;"
NL "    cvt.rn.f32.u32  p.x, x;"
NL "    fma.rn.f32      p.x, p.x, 1.0 / 768.0, -0.5;"
NL "    mul.f32         p.x, p.x, 2.0;"
NL "    cvt.rn.f32.u32  p.y, y;"
NL "    fma.rn.f32      p.y, p.y, 1.0 / 768.0, -0.5;"
NL "    mul.f32         p.y, p.y, 2.0;"
        /* load triangle points */
NL "    .reg .v2 .f32   p0, p1, p2;"
NL "    ld.const.v2.f32 p0, [nPoint];"
NL "    ld.const.v2.f32 p1, [nPoint + 8];"
NL "    ld.const.v2.f32 p2, [nPoint + 16];"
        /* load rotation matrix */
NL "    .reg .v2 .f32   rotZ<2>;"
NL "    ld.const.v2.f32 rotZ0, [nRotZ];"
NL "    ld.const.v2.f32 rotZ1, [nRotZ + 8];"
        /* rotate triangle points */
        mul_m2v2_f32    ("p0", "rotZ", "p0")
        mul_m2v2_f32    ("p1", "rotZ", "p1")
        mul_m2v2_f32    ("p2", "rotZ", "p2")
        /* check if point is inside the triangle */
NL "    .reg .f32       sign0, sign1, sign2;"
NL "    call            (sign0), ComputeSign, (p1, p0, p);"
NL "    call            (sign1), ComputeSign, (p2, p1, p);"
NL "    call            (sign2), ComputeSign, (p0, p2, p);"
NL "    .reg .pred      q;"
NL "    setp.gt.f32     q, sign0, 0.0;"
NL "    setp.gt.and.f32 q, sign1, 0.0, q;"
NL "    setp.gt.and.f32 q, sign2, 0.0, q;"
        /* set output color */
NL "    .reg .v4 .f32   color;"
NL "    @!q mov.v4.f32  color, { 0.0, 0.0, 0.0, 1.0 };"
NL "    @q mov.v4.f32   color, { 0.0, 1.0, 0.0, 1.0 };"
        /* write output color */
NL "    mul.lo.u32              x, x, 16;"
NL "    sust.b.2d.v4.b32.trap   [nFb, { x, y }], color;"
NL "    ret;"
NL "}";

static GLuint nProgram[2] = { 0 };
static GLuint nVBO = 0;
static GLuint64 nVBOAddress = 0;
static GLuint nGLTexture = 0;

static CUcontext nCUContext = NULL;
static CUmodule nCUModule = NULL;
static CUfunction nCUFunction = NULL;
static CUgraphicsResource nCUTexture = NULL;
static CUsurfref nCUSurface = NULL;
static CUdeviceptr nRotZPtr = 0;

int ghAppBegin() {
    CUdevice device;

    if (!GLEW_NV_gpu_program5) {
        GH_ERROR_SET("Program requires GL_NV_gpu_program5 extension.");
        return -1;
    }

    /* create programs */
    glGenProgramsARB(2, nProgram);

    glBindProgramARB(GL_VERTEX_PROGRAM_ARB, nProgram[0]);
    if (ghNVProgramSetString(GL_VERTEX_PROGRAM_ARB, nVertexProgramSrc) < 0) {
        return -1;
    }

    glBindProgramARB(GL_FRAGMENT_PROGRAM_ARB, nProgram[1]);
    if (ghNVProgramSetString(GL_FRAGMENT_PROGRAM_ARB, nFragmentProgramSrc) < 0) {
        return -1;
    }

    /* create vbo */ {
        float vertices[] = { -1, -1, -1, 1, 1, -1, 1, 1 };
        glGenBuffers(1, &nVBO);
        glBindBuffer(GL_ARRAY_BUFFER, nVBO);
        glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);
        glGetBufferParameterui64vNV(GL_ARRAY_BUFFER, GL_BUFFER_GPU_ADDRESS_NV, &nVBOAddress);
        glMakeBufferResidentNV(GL_ARRAY_BUFFER, GL_READ_ONLY);
    }

    /* create texture */
    glGenTextures(1, &nGLTexture);
    glBindTexture(GL_TEXTURE_RECTANGLE, nGLTexture);
    glTexParameteri(GL_TEXTURE_RECTANGLE, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_RECTANGLE, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexImage2D(GL_TEXTURE_RECTANGLE, 0, GL_RGBA32F, GH_APP_VIEWPORT_WIDTH, GH_APP_VIEWPORT_HEIGHT, 0, GL_RGBA, GL_FLOAT, NULL);

    /* set GL state */
    glEnable(GL_VERTEX_PROGRAM_ARB);
    glEnable(GL_FRAGMENT_PROGRAM_ARB);

    /* setup CUDA */
    if (cuInit(0) != CUDA_SUCCESS) {
        GH_ERROR_SET("cuInit failed.");
        return -1;
    }

    cuDeviceGet(&device, 0);
    if (cuGLCtxCreate(&nCUContext, 0, device) != CUDA_SUCCESS) {
        GH_ERROR_SET("cuCtxCreate failed.");
        return -1;
    }

    if (ghCUModuleLoad(&nCUModule, nTriModuleSrc) < 0) {
        return -1;
    }

    if (cuModuleGetFunction(&nCUFunction, nCUModule, "Draw") != CUDA_SUCCESS) {
        GH_ERROR_SET("cuModuleGetFunction failed.");
        return -1;
    }

    if (cuGraphicsGLRegisterImage(&nCUTexture, nGLTexture, GL_TEXTURE_RECTANGLE,
        CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD | CU_GRAPHICS_REGISTER_FLAGS_SURFACE_LDST) != CUDA_SUCCESS) {
        GH_ERROR_SET("cuGraphicsGLRegisterBuffer failed.");
        return -1;
    }

    if (cuModuleGetSurfRef(&nCUSurface, nCUModule, "nFb") != CUDA_SUCCESS) {
        GH_ERROR_SET("cuModuleGetSurfRef failed.");
        return -1;
    }

    if (cuModuleGetGlobal(&nRotZPtr, NULL, nCUModule, "nRotZ") != CUDA_SUCCESS) {
        GH_ERROR_SET("cuModuleGetGlobal failed.");
        return -1;
    }

    return 0;
}

int ghAppFrame(double dt, double time) {
    CUarray arr;

    /* compute rotation matrix */
    float sinv = sinf((float) time);
    float cosv = cosf((float) time);
    float rotZ[4] = {
        cosv, -sinv,
        sinv, cosv
    };

    cuMemcpyHtoD(nRotZPtr, rotZ, sizeof(rotZ));

    /* launch CUDA kernel */
    cuGraphicsMapResources(1, &nCUTexture, NULL);
    cuGraphicsSubResourceGetMappedArray(&arr, nCUTexture, 0, 0);
    cuSurfRefSetArray(nCUSurface, arr, 0);
    cuLaunchKernel(nCUFunction, 48, 48, 1, 16, 16, 1, 0, NULL, NULL, NULL);
    cuGraphicsUnmapResources(1, &nCUTexture, NULL);

    /* draw */
    glBindProgramARB(GL_VERTEX_PROGRAM_ARB, nProgram[0]);
    glBindProgramARB(GL_FRAGMENT_PROGRAM_ARB, nProgram[1]);
    glBindMultiTextureEXT(GL_TEXTURE0, GL_TEXTURE_RECTANGLE, nGLTexture);
    glEnableClientState(GL_VERTEX_ATTRIB_ARRAY_UNIFIED_NV);
    glEnableVertexAttribArray(0);
    glVertexAttribFormatNV(0, 2, GL_FLOAT, GL_FALSE, 2 * sizeof(float));
    glBufferAddressRangeNV(GL_VERTEX_ATTRIB_ARRAY_ADDRESS_NV, 0, nVBOAddress, 4 * 2 * sizeof(float));
    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

    return 0;
}

int ghAppEnd() {
    if (nCUTexture) cuGraphicsUnregisterResource(nCUTexture);
    if (nCUModule) cuModuleUnload(nCUModule);
    if (nCUContext) cuCtxDestroy(nCUContext);

    glDeleteBuffers(1, &nVBO);
    glDeleteProgramsARB(2, nProgram);
    glDeleteTextures(1, &nGLTexture);

    return 0;
}

int ghAppMouseCallback(int x, int y, int action) {
    return 0;
}
